from typing import List

import pandas as pd
import pyomo
import pyomo.environ as pe
from pyomo.opt import SolverStatus, TerminationCondition

from src.generator import Generator


def run_market_clearing(generators: List[Generator],
                        loads: pd.Series,
                        time_steps: []
                        ):
    ################################################################################
    # Initiate pyomo model
    ################################################################################
    m = pe.ConcreteModel()
    # Make duals available
    m.dual = pe.Suffix(direction=pe.Suffix.IMPORT)
    # Set of all generators
    m.generators = pe.Set(initialize=generators, dimen=1)
    m.loads = loads
    m.time_steps = pe.Set(initialize=time_steps, dimen=1)

    ################################################################################
    # Variables
    ################################################################################
    # Production of each generator (>= 0)
    m.production = pe.Var(m.generators, m.time_steps, domain=pe.NonNegativeReals)

    ################################################################################
    # Objective
    # Total production cost
    ################################################################################
    def Total_Social_Cost(m):
        return (
            + sum(g.c2 * m.production[g, t] * m.production[g, t]
                  + g.c1 * m.production[g, t] for g in m.generators for t in m.time_steps)
        )

    m.objective = pe.Objective(rule=Total_Social_Cost, sense=pe.minimize)

    ################################################################################
    # Constraints
    ################################################################################
    # Production capacity
    def Production_Upper_Bound_Rule(m, g, t):
        return m.production[g, t] <= g.capacity

    m.Production_Upper_Bound = pe.Constraint(m.generators, m.time_steps, rule=Production_Upper_Bound_Rule)

    # Market balance
    def Market_Balance_Rule(m, t):
        return sum(m.production[g, t] for g in m.generators) == m.loads[t]

    m.Market_Balance = pe.Constraint(m.time_steps, rule=Market_Balance_Rule)

    ################################################################################
    # solve the model
    ################################################################################
    solver = pyomo.opt.SolverFactory('gurobi')
    result = solver.solve(m, tee=True, keepfiles=False)  # tee: show solver info or not

    if result.solver.termination_condition == TerminationCondition.optimal or \
            result.solver.termination_condition == TerminationCondition.locallyOptimal:
        # m.display()  # display all the results
        print("Model has been solved.")
        return m
    else:
        print("Something went wrong.")
        return None
